﻿#ifndef __ROAD_LINE_FAR_DIRECTION_MODEL_H__
#define __ROAD_LINE_FAR_DIRECTION_MODEL_H__

#include "RoadLineUtils.h"
#include "JsColorModel.h"
#include "base64.h"
#include "models/SoftMaxModel.h"

using SoftMaxModel = darknet::SoftMaxModel;
class RoadLineFarDirectionModel : public SoftMaxModel
{
    using Options   = darknet::Options;
    using DataImage = darknet::DataImage;

public:
    static RoadLineFarDirectionModel *GetInstance()
    {
        static RoadLineFarDirectionModel s_instance;
        return &s_instance;
    }

public:
    RoadLineFarDirectionModel()
    {
        m_max_epochs      = 2000;
        m_batch           = 4;
        m_auto_stop_train = false;
        m_jitter          = 0.f;
        m_distort         = true;
        m_augment         = true;
        m_model_name      = "roadlines_far_direction";

#ifdef __EMSCRIPTEN__
        std::string root_directory = "/assets/" + GetModelName();
#else
        std::string root_directory = "D:/workspace/cpp/darknet/web/asset/assets/" + GetModelName();
#endif
        std::string model_directory   = root_directory + "/models";
        std::string model_filepath    = model_directory + "/" + GetDefaultModelFileName();
        std::string dataset_directory = root_directory + "/datasets";
        std::string dataset_filepath  = dataset_directory + "/dataset.txt";

        SetRootDirectory(root_directory);
        SetModelDirectory(model_directory);
        SetModelPath(model_filepath);
        SetDataSetDirectory(dataset_directory);
        SetDataSetPath(dataset_filepath);
#ifdef __EMSCRIPTEN__
        if (std::filesystem::exists(model_filepath) && LoadModelFromFile(model_filepath))
        {
            std::cout << "load prebuilt model from " << model_filepath << std::endl;
        }
        LoadDataSet();
        // std::cout << "load default dataset: " << GetDataSetPath() << " size = " << GetDataSet()->size() << std::endl;
#else
        std::filesystem::create_directories(root_directory);
        std::filesystem::create_directories(model_directory);
        std::filesystem::create_directories(dataset_directory);
#endif

        CreateGetClassifyIndex(reinterpret_cast<const char*>(u8R"(直行)"));
        CreateGetClassifyIndex(reinterpret_cast<const char*>(u8R"(直行左转15度)"));
        CreateGetClassifyIndex(reinterpret_cast<const char*>(u8R"(直行右转15度)"));
        CreateGetClassifyIndex(reinterpret_cast<const char*>(u8R"(直行左转30度)"));
        CreateGetClassifyIndex(reinterpret_cast<const char*>(u8R"(直行右转30度)"));
        CreateGetClassifyIndex(reinterpret_cast<const char*>(u8R"(左转15度)"));
        CreateGetClassifyIndex(reinterpret_cast<const char*>(u8R"(左转30度)"));
        CreateGetClassifyIndex(reinterpret_cast<const char*>(u8R"(左转45度)"));
        CreateGetClassifyIndex(reinterpret_cast<const char*>(u8R"(左转90度)"));
        CreateGetClassifyIndex(reinterpret_cast<const char*>(u8R"(右转15度)"));
        CreateGetClassifyIndex(reinterpret_cast<const char*>(u8R"(右转30度)"));
        CreateGetClassifyIndex(reinterpret_cast<const char*>(u8R"(右转45度)"));
        CreateGetClassifyIndex(reinterpret_cast<const char*>(u8R"(右转90度)"));
        CreateGetClassifyIndex(reinterpret_cast<const char*>(u8R"(禁止通行)"));
    }

    virtual void LoadNet()
    {
        if (GetNet()->IsLoaded()) return;

        auto net     = GetNet();
        auto options = std::make_shared<std::vector<std::shared_ptr<Options>>>();
        auto option  = std::make_shared<Options>("net");
        option->SetOption("width", RoadLineUtils::GetNetWidth());
        option->SetOption("height", RoadLineUtils::GetNetHeight());
        option->SetOption("channels", 3);
        options->push_back(option);

        option = std::make_shared<Options>("convolutional");
        option->SetOption("size", 3);
        option->SetOption("stride", 2);
        option->SetOption("filters", 16);
        options->push_back(option);

        option = std::make_shared<Options>("convolutional");
        option->SetOption("size", 3);
        option->SetOption("stride", 2);
        option->SetOption("filters", 32);
        options->push_back(option);
        
        option = std::make_shared<Options>("convolutional");
        option->SetOption("size", 3);
        option->SetOption("stride", 2);
        option->SetOption("filters", 32);
        options->push_back(option);

        option = std::make_shared<Options>("convolutional");
        option->SetOption("size", 3);
        option->SetOption("stride", 2);
        option->SetOption("filters", 16);
        options->push_back(option);

        option = std::make_shared<Options>("connected");
        option->SetOption("output", 1);
        option->SetOption("activation", "linear");
        options->push_back(option);

        option = std::make_shared<Options>("softmax");
        options->push_back(option);

        net->Load(options);
        net->SetBatch(GetBatch());
    }

    virtual void LoadDataSet()
    {
        SoftMaxModel::LoadDataSet();
    }

    virtual std::shared_ptr<DataImage> HandleInputImage(std::shared_ptr<DataImage> input_image)
    {
        return input_image->Scale(RoadLineUtils::GetNetImageScale(), RoadLineUtils::GetNetImageScale());
    }
};

#endif